"""
Kafka客户端模块

本模块提供Kafka生产者和消费者实现，用于处理空气质量数据的消息传递。
主要功能:
1. 发送原始空气质量数据到Kafka
2. 发送预测结果到Kafka 
3. 消费原始数据用于实时预测
4. 批量消费数据用于离线处理

使用说明:
1. 生产者: 创建AirQualityProducer实例，调用send_raw_data/send_prediction
2. 消费者: 创建AirQualityConsumer实例，调用consume_messages/consume_batch
3. 配置: 通过kafka_config模块获取配置
"""

import json
import logging
from kafka import KafkaProducer, KafkaConsumer
from kafka.errors import KafkaError
from datetime import datetime
from config.kafka_config import get_kafka_config

# 配置日志
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)
logger = logging.getLogger('kafka_client')

class AirQualityProducer:
    """空气质量数据生产者"""
    
    def __init__(self, config=None):
        """
        初始化Kafka生产者
        
        参数:
        config: Kafka配置，默认使用kafka_config中的配置
        """
        if config is None:
            config = get_kafka_config()
            
        self.producer = KafkaProducer(
            bootstrap_servers=config['bootstrap_servers'],
            value_serializer=lambda v: json.dumps(v).encode('utf-8'),
            key_serializer=lambda k: k.encode('utf-8') if k else None,
            acks='all'  # 确保消息被所有副本确认
        )
        self.raw_data_topic = config['raw_data_topic']
        self.prediction_topic = config['prediction_topic']
        
    def send_raw_data(self, data, callback=None):
        """
        发送原始空气质量数据
        
        参数:
            data (dict): 要发送的空气质量数据，应包含污染物数值
            callback (function): 可选，发送完成后的回调函数
            
        返回:
            None
            
        示例:
            producer.send_raw_data({
                'PM2_5': 35.2,
                'PM10': 78.5,
                'timestamp': '2025-04-09T21:00:00'
            })
        """
        # 添加时间戳作为消息键，确保同一时间的数据被分配到同一分区
        timestamp = data.get('timestamp', datetime.now().isoformat())
        
        # 发送消息
        future = self.producer.send(
            self.raw_data_topic,
            key=timestamp,
            value=data
        )
        
        # 如果提供了回调，则应用它
        if callback:
            future.add_callback(callback)
        else:
            future.add_callback(self._on_send_success)
            future.add_errback(self._on_send_error)
        
        # 确保所有消息被发送
        self.producer.flush()
    
    def send_prediction(self, prediction_data, callback=None):
        """
        发送预测结果
        
        参数:
            prediction_data (dict): 预测结果数据，应包含预测值和置信度
            callback (function): 可选，发送完成后的回调函数
            
        返回:
            None
            
        示例:
            producer.send_prediction({
                'predicted_PM2_5': 42.1,
                'confidence': 0.85,
                'timestamp': '2025-04-09T22:00:00'
            })
        """
        # 添加时间戳
        timestamp = prediction_data.get('timestamp', datetime.now().isoformat())
        
        # 发送消息
        future = self.producer.send(
            self.prediction_topic,
            key=timestamp,
            value=prediction_data
        )
        
        # 如果提供了回调，则应用它
        if callback:
            future.add_callback(callback)
        else:
            future.add_callback(self._on_send_success)
            future.add_errback(self._on_send_error)
        
        # 确保所有消息被发送
        self.producer.flush()
    
    def _on_send_success(self, record_metadata):
        """发送成功的回调"""
        logger.info(f"消息成功发送到 {record_metadata.topic} 分区 {record_metadata.partition} 偏移量 {record_metadata.offset}")
    
    def _on_send_error(self, exc):
        """发送失败的回调"""
        logger.error(f"消息发送失败: {exc}")
    
    def close(self):
        """关闭生产者"""
        self.producer.close()


class AirQualityConsumer:
    """空气质量数据消费者"""
    
    def __init__(self, topics=None, config=None):
        """
        初始化Kafka消费者
        
        参数:
        topics: 要消费的主题列表，默认为raw_data_topic
        config: Kafka配置，默认使用kafka_config中的配置
        """
        if config is None:
            config = get_kafka_config()
            
        if topics is None:
            topics = [config['raw_data_topic']]
        
        self.consumer = KafkaConsumer(
            *topics,
            bootstrap_servers=config['bootstrap_servers'],
            group_id=config['group_id'],
            auto_offset_reset=config['auto_offset_reset'],
            enable_auto_commit=config['enable_auto_commit'],
            auto_commit_interval_ms=config['auto_commit_interval_ms'],
            value_deserializer=lambda x: json.loads(x.decode('utf-8'))
        )
    
    def consume_messages(self, callback=None, timeout_ms=1000):
        """
        消费消息
        
        参数:
            callback (function): 处理消息的回调函数，接收message对象
            timeout_ms (int): 轮询超时时间(毫秒)，默认1000ms
            
        返回:
            None
            
        注意:
            如果没有提供callback，默认会记录消息到日志
        """
        messages = self.consumer.poll(timeout_ms=timeout_ms)
        
        # 处理消息
        for topic_partition, records in messages.items():
            for record in records:
                if callback:
                    callback(record)
                else:
                    logger.info(f"接收到消息: Topic={record.topic}, Partition={record.partition}, "
                               f"Offset={record.offset}, Key={record.key}, Value={record.value}")
    
    def consume_batch(self, max_records=10, timeout_ms=1000):
        """
        批量消费消息
        
        参数:
        max_records: 最大记录数
        timeout_ms: 轮询超时时间（毫秒）
        
        返回:
        收集到的消息列表
        """
        batch = []
        start_time = datetime.now()
        timeout_delta = timeout_ms / 1000  # 转换为秒
        
        try:
            while len(batch) < max_records:
                # 检查是否超时
                if (datetime.now() - start_time).total_seconds() > timeout_delta:
                    break
                
                # 轮询一条消息
                records = self.consumer.poll(timeout_ms=100, max_records=1)
                
                for topic_partition, messages in records.items():
                    for message in messages:
                        batch.append(message.value)
                        # 如果达到目标数量，提前退出
                        if len(batch) >= max_records:
                            break
                    if len(batch) >= max_records:
                        break
        
        except Exception as e:
            logger.error(f"批量消费消息时出错: {str(e)}")
        
        return batch
    
    def subscribe(self, topics):
        """
        订阅主题
        
        参数:
        topics: 要订阅的主题列表
        """
        self.consumer.subscribe(topics)
    
    def close(self):
        """关闭消费者"""
        self.consumer.close()


# 辅助函数
def format_air_quality_message(data):
    """
    格式化空气质量数据为Kafka消息
    
    参数:
    data: 空气质量数据字典
    
    返回:
    格式化后的消息
    """
    # 确保有时间戳
    if 'timestamp' not in data:
        data['timestamp'] = datetime.now().isoformat()
    
    # 确保所有数值为浮点数
    for key in ['PM2_5', 'PM10', 'SO2', 'NO2', 'CO', 'O3', 'AQI']:
        if key in data and data[key] is not None:
            data[key] = float(data[key])
    
    return data
